"""
数据生成器 - 生成频谱和对应的标签
用于训练Transformer模型进行频谱分析
"""

import numpy as np
import json
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, asdict
from signal_generator import (
    SignalParams, HarmonicGenerator, ModulationGenerator,
    BearingFaultGenerator, CompositeSignalGenerator
)


@dataclass
class SignalLabel:
    """信号标签 - 包含所有需要识别的特征"""
    
    # 谐波相关
    has_harmonic: bool = False  # 是否有谐波
    harmonic_clusters: List[Dict] = None  # 谐波簇列表
    # 每个簇: {'base_freq': 50.0, 'num_harmonics': 5, 'amplitudes': [1.0, 0.5, ...]}
    
    # 调制相关
    has_modulation: bool = False  # 是否有调制
    modulation_info: List[Dict] = None  # 调制信息列表
    # 每个调制: {'carrier_freq': 500.0, 'modulation_freq': 30.0, 'num_sidebands': 1}
    
    # 轴承故障相关
    has_bearing_fault: bool = False  # 是否有轴承故障
    bearing_fault_info: Optional[Dict] = None
    # {'type': 'outer_race', 'fault_freq': 90.0, 'resonance_freq': 3000.0}
    
    def __post_init__(self):
        if self.harmonic_clusters is None:
            self.harmonic_clusters = []
        if self.modulation_info is None:
            self.modulation_info = []
    
    def to_dict(self) -> Dict:
        """转换为字典"""
        return asdict(self)
    
    def to_json(self) -> str:
        """转换为JSON字符串"""
        return json.dumps(self.to_dict())


class SpectrumDataGenerator:
    """频谱数据生成器"""
    
    def __init__(self, 
                 fs: float = 5120.0,
                 n_samples: int = 16384,
                 seed: Optional[int] = None):
        """
        初始化数据生成器
        
        参数:
            fs: 采样频率
            n_samples: 采样点数
            seed: 随机种子
        """
        self.params = SignalParams(fs=fs, n_samples=n_samples)
        if seed is not None:
            np.random.seed(seed)
        
        # 频率范围设置
        self.max_freq = fs / 2  # Nyquist频率
        self.harmonic_base_freq_range = (5.0, 200.0)  # 谐波基频范围
        self.modulation_carrier_range = (200.0, 1500.0)  # 载波频率范围
        self.modulation_freq_range = (5.0, 100.0)  # 调制频率范围
        
    def compute_spectrum(self, signal: np.ndarray) -> np.ndarray:
        """
        计算信号的单边频谱
        
        返回:
            spectrum: 频谱幅值 (长度为n_samples//2)
        """
        fft_vals = np.fft.fft(signal)
        # 只取正频率部分
        spectrum = np.abs(fft_vals[:len(fft_vals)//2]) / len(signal)
        # 除了直流分量，其他频率分量乘以2
        spectrum[1:] *= 2
        return spectrum
    
    def get_freq_axis(self) -> np.ndarray:
        """获取频率轴"""
        return np.fft.fftfreq(self.params.n_samples, 1/self.params.fs)[:self.params.n_samples//2]
    
    def generate_pure_harmonic(self) -> Tuple[np.ndarray, SignalLabel]:
        """生成纯谐波信号"""
        gen = HarmonicGenerator(self.params)
        
        # 随机生成1-3个谐波簇
        num_clusters = np.random.randint(1, 4)
        clusters_info = []
        combined_signal = np.zeros(self.params.n_samples)
        
        for _ in range(num_clusters):
            # 随机基频
            base_freq = np.random.uniform(*self.harmonic_base_freq_range)
            # 随机谐波数量 (2-6)
            num_harmonics = np.random.randint(2, 7)
            
            # 随机幅值（递减）
            amplitudes = [np.random.uniform(0.5, 1.5) / (i + 1) for i in range(num_harmonics)]
            
            signal, _ = gen.generate(
                fundamental_freq=base_freq,
                num_harmonics=num_harmonics,
                amplitudes=amplitudes
            )
            
            combined_signal += signal
            
            clusters_info.append({
                'base_freq': float(base_freq),
                'num_harmonics': num_harmonics,
                'amplitudes': [float(a) for a in amplitudes]
            })
        
        # 添加噪声
        noise_level = np.random.uniform(0.01, 0.1)
        combined_signal += np.random.normal(0, noise_level, len(combined_signal))
        
        label = SignalLabel(
            has_harmonic=True,
            harmonic_clusters=clusters_info
        )
        
        spectrum = self.compute_spectrum(combined_signal)
        return spectrum, label
    
    def generate_modulation_signal(self) -> Tuple[np.ndarray, SignalLabel]:
        """生成调制信号"""
        gen = ModulationGenerator(self.params)
        
        # 随机生成1-2个调制成分
        num_modulations = np.random.randint(1, 3)
        modulation_infos = []
        combined_signal = np.zeros(self.params.n_samples)
        
        for _ in range(num_modulations):
            carrier_freq = np.random.uniform(*self.modulation_carrier_range)
            modulation_freq = np.random.uniform(*self.modulation_freq_range)
            modulation_index = np.random.uniform(0.3, 0.9)
            
            signal, _ = gen.generate_am(
                carrier_freq=carrier_freq,
                modulation_freq=modulation_freq,
                carrier_amp=np.random.uniform(0.5, 1.5),
                modulation_index=modulation_index
            )
            
            combined_signal += signal
            
            # 边带数量（AM信号有2个一阶边带）
            modulation_infos.append({
                'carrier_freq': float(carrier_freq),
                'modulation_freq': float(modulation_freq),
                'num_sidebands': 1,  # 一阶边带（左右各一个）
                'sideband_spacing': float(modulation_freq)
            })
        
        # 添加噪声
        noise_level = np.random.uniform(0.01, 0.1)
        combined_signal += np.random.normal(0, noise_level, len(combined_signal))
        
        label = SignalLabel(
            has_modulation=True,
            modulation_info=modulation_infos
        )
        
        spectrum = self.compute_spectrum(combined_signal)
        return spectrum, label
    
    def generate_harmonic_with_modulation(self) -> Tuple[np.ndarray, SignalLabel]:
        """生成谐波+调制的复合信号"""
        # 生成谐波
        harmonic_gen = HarmonicGenerator(self.params)
        base_freq = np.random.uniform(*self.harmonic_base_freq_range)
        num_harmonics = np.random.randint(2, 5)
        amplitudes = [np.random.uniform(0.5, 1.0) / (i + 1) for i in range(num_harmonics)]
        
        harmonic_signal, _ = harmonic_gen.generate(
            fundamental_freq=base_freq,
            num_harmonics=num_harmonics,
            amplitudes=amplitudes
        )
        
        harmonic_cluster = {
            'base_freq': float(base_freq),
            'num_harmonics': num_harmonics,
            'amplitudes': [float(a) for a in amplitudes]
        }
        
        # 生成调制
        mod_gen = ModulationGenerator(self.params)
        carrier_freq = np.random.uniform(*self.modulation_carrier_range)
        modulation_freq = np.random.uniform(*self.modulation_freq_range)
        
        mod_signal, _ = mod_gen.generate_am(
            carrier_freq=carrier_freq,
            modulation_freq=modulation_freq,
            carrier_amp=np.random.uniform(0.5, 1.0),
            modulation_index=np.random.uniform(0.4, 0.8)
        )
        
        modulation_info_item = {
            'carrier_freq': float(carrier_freq),
            'modulation_freq': float(modulation_freq),
            'num_sidebands': 1,
            'sideband_spacing': float(modulation_freq)
        }
        
        # 叠加
        combined_signal = harmonic_signal + mod_signal
        
        # 添加噪声
        noise_level = np.random.uniform(0.02, 0.1)
        combined_signal += np.random.normal(0, noise_level, len(combined_signal))
        
        label = SignalLabel(
            has_harmonic=True,
            harmonic_clusters=[harmonic_cluster],
            has_modulation=True,
            modulation_info=[modulation_info_item]
        )
        
        spectrum = self.compute_spectrum(combined_signal)
        return spectrum, label
    
    def generate_bearing_fault(self) -> Tuple[np.ndarray, SignalLabel]:
        """生成轴承故障信号"""
        gen = BearingFaultGenerator(self.params)
        
        rotation_freq = np.random.uniform(10.0, 50.0)
        fault_type = np.random.choice(['outer_race', 'inner_race'])
        
        if fault_type == 'outer_race':
            signal, info = gen.generate_outer_race_fault(
                rotation_freq=rotation_freq,
                num_balls=np.random.randint(7, 12),
                impact_amplitude=np.random.uniform(1.0, 2.5),
                resonance_freq=np.random.uniform(2000.0, 2500.0)
            )
            fault_freq = info['BPFO']
        else:
            signal, info = gen.generate_inner_race_fault(
                rotation_freq=rotation_freq,
                num_balls=np.random.randint(7, 12),
                impact_amplitude=np.random.uniform(1.0, 2.5),
                resonance_freq=np.random.uniform(2000.0, 2500.0)
            )
            fault_freq = info['BPFI']
        
        bearing_info = {
            'type': fault_type,
            'fault_freq': float(fault_freq),
            'resonance_freq': float(info['resonance_freq'])
        }
        
        # 添加噪声
        noise_level = np.random.uniform(0.03, 0.15)
        signal += np.random.normal(0, noise_level, len(signal))
        
        label = SignalLabel(
            has_bearing_fault=True,
            bearing_fault_info=bearing_info
        )
        
        spectrum = self.compute_spectrum(signal)
        return spectrum, label
    
    def generate_complex_signal(self) -> Tuple[np.ndarray, SignalLabel]:
        """生成复杂信号（包含多种成分）"""
        composite_gen = CompositeSignalGenerator(self.params)
        
        # 随机决定包含哪些成分
        include_harmonic = np.random.random() > 0.3
        include_modulation = np.random.random() > 0.3
        include_bearing = np.random.random() > 0.5
        
        signals = []
        weights = []
        
        harmonic_clusters = []
        modulation_infos = []
        bearing_info = None
        
        # 谐波
        if include_harmonic:
            harmonic_gen = HarmonicGenerator(self.params)
            base_freq = np.random.uniform(*self.harmonic_base_freq_range)
            num_harmonics = np.random.randint(2, 5)
            amplitudes = [np.random.uniform(0.5, 1.0) / (i + 1) for i in range(num_harmonics)]
            
            signal, _ = harmonic_gen.generate(
                fundamental_freq=base_freq,
                num_harmonics=num_harmonics,
                amplitudes=amplitudes
            )
            signals.append(signal)
            weights.append(np.random.uniform(0.8, 1.2))
            
            harmonic_clusters.append({
                'base_freq': float(base_freq),
                'num_harmonics': num_harmonics,
                'amplitudes': [float(a) for a in amplitudes]
            })
        
        # 调制
        if include_modulation:
            mod_gen = ModulationGenerator(self.params)
            carrier_freq = np.random.uniform(*self.modulation_carrier_range)
            modulation_freq = np.random.uniform(*self.modulation_freq_range)
            
            signal, _ = mod_gen.generate_am(
                carrier_freq=carrier_freq,
                modulation_freq=modulation_freq,
                carrier_amp=np.random.uniform(0.5, 1.0),
                modulation_index=np.random.uniform(0.4, 0.8)
            )
            signals.append(signal)
            weights.append(np.random.uniform(0.6, 1.0))
            
            modulation_infos.append({
                'carrier_freq': float(carrier_freq),
                'modulation_freq': float(modulation_freq),
                'num_sidebands': 1,
                'sideband_spacing': float(modulation_freq)
            })
        
        # 轴承故障
        if include_bearing:
            bearing_gen = BearingFaultGenerator(self.params)
            rotation_freq = np.random.uniform(10.0, 40.0)
            
            signal, info = bearing_gen.generate_outer_race_fault(
                rotation_freq=rotation_freq,
                num_balls=np.random.randint(7, 11),
                impact_amplitude=np.random.uniform(1.0, 2.0),
                resonance_freq=np.random.uniform(2000.0, 2500.0)
            )
            signals.append(signal)
            weights.append(np.random.uniform(1.0, 1.5))
            
            bearing_info = {
                'type': 'outer_race',
                'fault_freq': float(info['BPFO']),
                'resonance_freq': float(info['resonance_freq'])
            }
        
        # 如果没有任何成分，至少生成一个谐波
        if len(signals) == 0:
            return self.generate_pure_harmonic()
        
        # 叠加
        combined_signal = composite_gen.combine_signals(signals, weights)
        
        # 添加噪声
        noise_level = np.random.uniform(0.05, 0.15)
        combined_signal += np.random.normal(0, noise_level, len(combined_signal))
        
        label = SignalLabel(
            has_harmonic=include_harmonic,
            harmonic_clusters=harmonic_clusters,
            has_modulation=include_modulation,
            modulation_info=modulation_infos,
            has_bearing_fault=include_bearing,
            bearing_fault_info=bearing_info
        )
        
        spectrum = self.compute_spectrum(combined_signal)
        return spectrum, label
    
    def generate_batch(self, 
                      batch_size: int,
                      signal_types: List[str] = None) -> Tuple[np.ndarray, List[SignalLabel]]:
        """
        生成一批数据
        
        参数:
            batch_size: 批次大小
            signal_types: 信号类型列表，如果为None则随机选择
                         可选: 'harmonic', 'modulation', 'harmonic_modulation', 
                              'bearing', 'complex'
        
        返回:
            spectra: 频谱数组 (batch_size, n_samples//2)
            labels: 标签列表
        """
        if signal_types is None:
            signal_types = ['harmonic', 'modulation', 'harmonic_modulation', 
                          'bearing', 'complex']
        
        spectra = []
        labels = []
        
        for _ in range(batch_size):
            sig_type = np.random.choice(signal_types)
            
            if sig_type == 'harmonic':
                spectrum, label = self.generate_pure_harmonic()
            elif sig_type == 'modulation':
                spectrum, label = self.generate_modulation_signal()
            elif sig_type == 'harmonic_modulation':
                spectrum, label = self.generate_harmonic_with_modulation()
            elif sig_type == 'bearing':
                spectrum, label = self.generate_bearing_fault()
            elif sig_type == 'complex':
                spectrum, label = self.generate_complex_signal()
            else:
                raise ValueError(f"Unknown signal type: {sig_type}")
            
            spectra.append(spectrum)
            labels.append(label)
        
        return np.array(spectra), labels


def save_dataset(filename: str, spectra: np.ndarray, labels: List[SignalLabel]):
    """保存数据集"""
    data = {
        'spectra': spectra.tolist(),
        'labels': [label.to_dict() for label in labels]
    }
    np.savez_compressed(filename, **data)
    print(f"数据集已保存到 {filename}")


def load_dataset(filename: str) -> Tuple[np.ndarray, List[SignalLabel]]:
    """加载数据集"""
    data = np.load(filename, allow_pickle=True)
    spectra = np.array(data['spectra'])
    labels = [SignalLabel(**label_dict) for label_dict in data['labels']]
    return spectra, labels


if __name__ == "__main__":
    # 测试数据生成器
    print("测试数据生成器...")
    
    generator = SpectrumDataGenerator(seed=42)
    
    # 生成测试样本
    print("\n生成纯谐波信号...")
    spectrum, label = generator.generate_pure_harmonic()
    print(f"频谱形状: {spectrum.shape}")
    print(f"标签: {label.to_dict()}")
    
    print("\n生成调制信号...")
    spectrum, label = generator.generate_modulation_signal()
    print(f"标签: {label.to_dict()}")
    
    print("\n生成复合信号...")
    spectrum, label = generator.generate_complex_signal()
    print(f"标签: {label.to_dict()}")
    
    print("\n生成批量数据...")
    spectra, labels = generator.generate_batch(5)
    print(f"批量频谱形状: {spectra.shape}")
    print(f"批量大小: {len(labels)}")

